package typetcp

import (
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/fatih/color"
)

var (
	Ok = func(args ...interface{}) {
		L("[+] ", args...)
	}
	Failed = func(args ...interface{}) {
		L("[!] ", args...)
	}
	DEFAULT_TIMEOUT = 17
	wg              sync.WaitGroup
	END             = "[BREAK]"
	NeedRetry       = make(chan string)
)

type User struct {
	Name  string
	Pwd   string
	Other string
}

type BrutePro interface {
	Brute() bool
}

type Bruter struct {
	Targets        []string
	Users          []string
	Pwd            []string
	allCount       int
	finished       int
	finishedTarget map[string]bool
	jobs           chan string
	results        chan string
	doing          chan int
	sigals         chan int
}

type Auth struct {
	Type   string
	Target string
	User   string
	Pwd    string
	Extend string
}

func Lines(f string) (o []string) {
	buf, err := ioutil.ReadFile(f)
	if err != nil {
		return
		// log.Fatal(err)
	}
	for _, t := range strings.Split(string(buf), "\n") {
		if strings.HasPrefix(t, "#") {
			continue
		}
		if strings.TrimSpace(t) != "" {
			o = append(o, strings.TrimSpace(t))
		}
	}
	return
}

func NewBruter(target string, user, pwd string) (brute *Bruter) {
	brute = &Bruter{
		Targets:        Lines(target),
		Users:          Lines(user),
		Pwd:            Lines(pwd),
		jobs:           make(chan string, 256),
		results:        make(chan string, 256),
		doing:          make(chan int, 256),
		sigals:         make(chan int, 10),
		finishedTarget: make(map[string]bool),
	}
	if target != "" && len(brute.Targets) == 0 {
		brute.Targets = []string{target}
	}
	if user != "" && len(brute.Users) == 0 {
		brute.Users = []string{user}
	} else if user == "" {
		brute.Users = []string{"root"}
	}

	if pwd != "" && len(brute.Pwd) == 0 {
		brute.Pwd = []string{pwd}
	}

	return
}

func (bruter *Bruter) working(id int, uri string) {
	defer func() {
		<-bruter.doing
	}()
	if bruter.AutoBrute(uri) {
		bruter.results <- "Good " + uri
		bruter.finished++
		lock.Lock()
		bruter.finishedTarget[strings.SplitN(uri, "@", 2)[0]] = true
		lock.Unlock()
		bruter.sigals <- id
	} else {
		bruter.results <- "bad " + uri
	}
}

func (bruter *Bruter) WorkerLoop() {
	c := 0
	// Ok("run woker loop")

	for iduri := range bruter.jobs {
		// fmt.Println("ready :", len(bruter.jobs))

		c++
		ids := strings.SplitN(iduri, ":", 2)
		id, err := strconv.Atoi(ids[0])
		if err != nil {
			log.Println("job error:", err)
			break
		}
		uri := ids[1]
		if uri == END {
			bruter.results <- END
			break
		}

		bruter.doing <- 1
		go bruter.working(id, uri)

		// fmt.Println(uri, len(jobs), "do:", len(doing))
	}
}

func (bruter *Bruter) WorkerResult() {
	fc := 0
	all := len(bruter.Targets) * len(bruter.Users) * len(bruter.Pwd)
	// Ok("run Result recve channel")
	for res := range bruter.results {
		fc++
		if res != END {
			if strings.Contains(res, "Good") {
				Ok(res, "            ")
				// goods <- res
			} else {
				if fc%3 == 0 {
					fmt.Printf("\r%d/%d %6d/%d Doing:%d | %s               \r", bruter.finished, bruter.allCount, fc, all, len(bruter.doing), res)
				}
				// Failed(res)
			}
		} else {
			break
		}
	}
}

func (bruter *Bruter) InitPool() {
	bruter.jobs = make(chan string, 500)
	bruter.results = make(chan string, 500)
	bruter.doing = make(chan int, 500)
}

func (bruter *Bruter) BruteAll() {
	// bruter.allCount = len(bruter.Targets) * len(bruter.Users) * len(bruter.Pwd)
	// progress := pb.New(len(bruter.Targets) * len(bruter.Users) * len(bruter.Pwd))
	// defer progress.Finish()
	// allTargets := len(bruter.Targets)
	finishedTargets := 0
	// Good := make(chan string)

	go bruter.WorkerLoop()
	go bruter.WorkerResult()
	for id, target := range bruter.Targets {
		// fmt.Println("t:", target)

		dynamicTask := func(id int, t string) {
			defer func() {
				defer wg.Done()
				finishedTargets++
				bruter.allCount--
			}()

			wg.Add(1)
			bruter.allCount++
			// c := 0
			Ok("task id:", id, t)
			for _, u := range bruter.Users {
				for _, p := range bruter.Pwd {
					if _, ok := bruter.finishedTarget[t]; ok {
						return
					}
					bruteURI := fmt.Sprintf("%s@%s:%s", t, u, p)

					bruter.jobs <- fmt.Sprintf("%d:%s", id, bruteURI)
					// select {
					// case sig := <-bruter.sigals:
					// 	if sig == id {
					// 		goto EndBreak
					// 	}
					// case bruter.jobs <- fmt.Sprintf("%d:%s", id, bruteURI):
					// }

					// c++
					// if bruter.AutoBrute(bruteURI) {
					// 	Ok("Good", t, u, p)
					// 	break
					// } else {
					// 	fmt.Printf("\r%d/%d %d/%d | %s               \r", finishedTargets, allTargets, c, all, t)
					// }
				}

			}

			// time.Sleep(10 * time.Second)
			// Ok("running:", len(bruter.doing))
			for {
				if len(bruter.doing) == 0 {
					// bruter.jobs <- END
					break
				}
			}
			// EndBreak:
		}

		// if len(bruter.Users)*len(bruter.Pwd) > 100 {
		// 	dynamicTask(target)
		// } else {
		go dynamicTask(id, target)
		// }

	}
	if len(NeedRetry) > 0 {
		Ok("retry loop")
		for retry := range NeedRetry {
			if _, ok := bruter.finishedTarget[strings.SplitN(retry, "@", 2)[0]]; ok {
				continue
			}
			bruter.jobs <- fmt.Sprintf("%d:%s", -1, retry)
			// if strings.Contains(retry, "dr:dark") {
			// fmt.Println("retry :", retry, "                   ")
			if len(NeedRetry) == 0 {
				break
			}
			// }
		}
	}
	time.Sleep(1 * time.Second)
	Ok("----------------- seprate ----------------- ")
	for {
		time.Sleep(4 * time.Second)
		if len(bruter.jobs) == 0 && len(bruter.doing) == 0 {
			bruter.jobs <- END
			break
		}
	}
	wg.Wait()
	// for g := range Good {
	// 	Ok(g)
	// }
}

func (bruter *Bruter) AutoBrute(s string) bool {
	// Ok(s)
	// defer func() {
	// bruter.finished++
	// fmt.Printf("\rFinished %d/%d \r", bruter.finished, bruter.allCount)
	// }()
	if !strings.Contains(s, "://") {
		s = "ssh://" + s
	}
	if strings.HasPrefix(s, "ssh://") {
		return Ssh(s).Brute()
		// progress.Write().P
	} else if strings.HasPrefix(s, "http://") {
		return BaseHTTP(s).Brute()
	} else if strings.HasPrefix(s, "https://") {
		return BaseHTTPS(s).Brute()
	} else if strings.HasPrefix(s, "ftp://") {
		return Ftp(s).Brute()
	}
	return false
}

func L(pre string, all ...interface{}) {
	pre = color.New(color.FgGreen, color.Bold).Sprint(pre)
	end := color.New(color.Bold).Sprintln(all...)
	fmt.Print(pre, end)
}

// pro://ip:port@user:pwd
func ParseStringToBrute(s string) (a *Auth, err error) {
	fs := strings.SplitN(s, "@", 2)
	if len(fs) != 2 {
		return nil, errors.New("must |  protcol://ip:port@user:pwd you are:" + s)
	}
	afs := strings.SplitN(fs[1], ":", 2)
	if len(afs) != 2 {
		return nil, errors.New("must |  user:pwd, you are:" + s)
	}
	protocl := strings.SplitN(fs[0], "://", 2)[0]
	addr := strings.SplitN(fs[0], "://", 2)[1]
	a = new(Auth)
	a.Type = protocl
	a.Target = addr
	a.User = afs[0]
	a.Pwd = afs[1]
	return
}
